Add generic interface for SSM inference by santhnm2 · Pull Request #5382 · NVIDIA/Megatron-LM

santhnm2 · 2026-06-16T22:59:01Z

I, the PR author, have personally reviewed every line of this PR.

What does this PR do ?

Makes the following refactoring changes (no functionality changes):

Factors out common decode / prefill split logic for linear attention dynamic inference into ssm_inference.py
Re-organizes Mamba2 dynamic inference code to use the generic ssm_decode and ssm_prefill interface functions
Quarantines the legacy static Mamba2 inference code into _static_decode and _static_prefill functions.

⚠️ For major changes (either in lines of code or in its impact), please make sure to first share a design doc with the team. If you're unsure what's the best way to do so, contact @NVIDIA/mcore-oncall.

Issue tracking

For PRs from open-source community contributors:

New features: a linked issue is required. Please open a feature request and reference it here before submitting the PR.
Small updates (bug fixes, minor improvements): a linked issue is recommended and will accelerate the PR review process.

Linked issue:

Contribution process

Pre-checks

I have added relevant unit tests
I have added relevant functional tests
I have added proper typing to my code Typing guidelines
I have added relevant documentation
I have run the autoformatter.sh on my PR

Code review

Feel free to message or comment @NVIDIA/mcore-oncall to help accelerate your merge into main. The less complex your PR is, the faster it will be approved and merged!

All PRs start as draft. If you open a non-draft PR, it will be automatically converted to draft.

Step 1: Mark PR as "Ready for Review"

When your PR is ready, click Ready for Review.
An oncall reviewer is auto-assigned and expert reviewers are notified based on your changes.
- Some PRs may jump straight to step 2. This is determined by .github/CODEOWNERS.

⚠️ Only mark as ready once merge-conflicts are resolved and the CI is passing.
Final Review might get declined if these requirements are not fulfilled.

Step 2: Final Review

For PRs that change megatron/core, once all expert reviewers have approved, the Final Review label is applied automatically and final reviewers are assigned.

For PRs outside megatron/core, this step is skipped.

Step 3: Approved

Once all required reviewers have approved, the Approved label is applied automatically.

Merge

Any member of mcore-engineers will be able to merge your PR.

copy-pr-bot · 2026-06-16T22:59:04Z

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

Signed-off-by: Keshav Santhanam <ksanthanam@nvidia.com>

santhnm2 · 2026-08-04T15:51:52Z

/ok to test 74becf4

shanmugamr1992 · 2026-08-05T00:17:21Z

NeMo-RL functional tests

mcore sha: eb01b689c4ddf1034db987a25736aee12552d2ea · baseline main sha: 3aee84c392e3a265c497991d0d07b27bfffb2490 · bridge sha: 2940a635731da68c5960a020f8ea2f678fa1e073 · nemo rl sha: 86774472e229b457314418fe3f25e8e844775596 · suite: L1_Functional_Tests_Megatron_4 · cluster: oci-hsg · updated 2026-08-05 23:25 UTC

Test	Status	Comment
`grpo_megatron_generation`	🛠️ fixed	Passes only with the Megatron-Bridge fix in NVIDIA-NeMo/Megatron-Bridge#5357 applied. Without it this test fails at import time on megatron-core main too, so the break is not caused by this PR.
`grpo_megatron_generation_non_colocated`	🛠️ fixed	Passes only with the Megatron-Bridge fix in NVIDIA-NeMo/Megatron-Bridge#5357 applied. Without it this test fails at import time on megatron-core main too, so the break is not caused by this PR.
`grpo_megatron_generation_async`	🛠️ fixed	Passes only with the Megatron-Bridge fix in NVIDIA-NeMo/Megatron-Bridge#5357 applied. Without it this test fails at import time on megatron-core main too, so the break is not caused by this PR.
`grpo_megatron_generation_colocated_async`	🛠️ fixed	Passes only with the Megatron-Bridge fix in NVIDIA-NeMo/Megatron-Bridge#5357 applied. Without it this test fails at import time on megatron-core main too, so the break is not caused by this PR.
`grpo_megatron_generation_async_gym`	🛠️ fixed	Passes only with the Megatron-Bridge fix in NVIDIA-NeMo/Megatron-Bridge#5357 applied. Without it this test fails at import time on megatron-core main too, so the break is not caused by this PR.
`grpo_megatron_generation_topp_topk`	⚠️ fail (pre-existing)	Not caused by this PR, and not the Bridge crash either. With NVIDIA-NeMo/Megatron-Bridge#5357 applied this test runs to completion and misses a metric threshold: max(train/token_mult_prob_error) = 1.078047513961792 against < 1.06. The value is bit-identical on this PR's head (`eb01b68`) and on megatron-core main (`3aee84c`), so it is deterministic and pre-existing, not run-to-run noise. Tracked as topp-topk-token-mult-prob-error-over-threshold; no fix raised yet.
`grpo_megatron_generation_colocated_gym`	🛠️ fixed	Passes only with the Megatron-Bridge fix in NVIDIA-NeMo/Megatron-Bridge#5357 applied. Without it this test fails at import time on megatron-core main too, so the break is not caused by this PR.

No failure in this suite is attributable to this PR.

L1 was fully blocked for every PR by a Megatron-Bridge break: Megatron-LM#5865 turned mcore_fsdp_adapter.FullyShardedDataParallel into a factory function, so Bridge's isinstance() checks against it raise TypeError. That is fixed in draft NVIDIA-NeMo/Megatron-Bridge#5357, and this run carried it — hence 🛠️ fixed rather than plain pass. Until that merges, these tests fail on main too.

The one remaining failure, grpo_megatron_generation_topp_topk, reproduces bit-identically on megatron-core main, so it is pre-existing as well. Tracked in NVIDIA-NeMo/RL#3365.

L2 has not been run yet.

_{Posted by the nemo-rl-testing-agent. Re-runs edit this comment in place.}

santhnm2 · 2026-08-05T16:55:15Z

/ok to test eb01b68

Signed-off-by: Keshav Santhanam <ksanthanam@nvidia.com>

santhnm2 · 2026-08-06T20:04:15Z

/ok to test 83e1849

santhnm2 · 2026-08-07T19:24:41Z

/ok to test 8884612

svcnvidia-nemo-ci · 2026-08-08T19:54:57Z

🔄 Merge queue validation started!

You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/31275604068

GDP's in_proj.weight is GTP-sliced along axis 0 with alignment padding; the [z|V*|K*|Q|b*|a] split boundaries don't line up with GTP slice boundaries, so sharded_state_dict all-gathers the shards back to TP-local size (stripping pad rows) and falls through to the non-GTP split path -- the saved checkpoint matches a non-GTP run. On load, the factory merge_fn re-pads and re-slices per gtp_remat rank, mirroring _gtp_slice_one_param. (squashed cherry-pick of unmerged PR NVIDIA#6503 @ 04fc214, adapted to the pre-NVIDIA#6443 GDP file: NVIDIA#6443/NVIDIA#5382 (dynamic-inference SSM interface) are not carried, so the import block keeps is_using_quantization_scales)

Merge with main

87f2871

Signed-off-by: Keshav Santhanam <ksanthanam@nvidia.com>

santhnm2 force-pushed the ssm_inference_interface branch from 562ab67 to 87f2871 Compare August 4, 2026 15:42

santhnm2 added 3 commits August 4, 2026 08:43

Restore ssm_inference.py

8a0037d

Signed-off-by: Keshav Santhanam <ksanthanam@nvidia.com>

Clean up comments

061b116

Signed-off-by: Keshav Santhanam <ksanthanam@nvidia.com>

Update comments

74becf4

Signed-off-by: Keshav Santhanam <ksanthanam@nvidia.com>

santhnm2 marked this pull request as ready for review August 4, 2026 15:51

santhnm2 requested review from a team as code owners August 4, 2026 15:51

nvidia-megatron-lm-release-bot Bot added the community-request label Aug 4, 2026

github-actions Bot added the complexity: high label Aug 4, 2026

copy-pr-bot Bot temporarily deployed to public August 4, 2026 15:54 Inactive

copy-pr-bot Bot temporarily deployed to test August 4, 2026 15:55 Inactive

copy-pr-bot Bot temporarily deployed to public August 4, 2026 15:58 Inactive

copy-pr-bot Bot temporarily deployed to public August 4, 2026 16:10 Inactive

shanmugamr1992 added the NemoRLTest label Aug 5, 2026

Merge branch 'main' into ssm_inference_interface

f05c3ee

kvareddy approved these changes Aug 5, 2026

View reviewed changes

santhnm2 added the Run functional tests label Aug 5, 2026

Merge branch 'main' into ssm_inference_interface

eb01b68

copy-pr-bot Bot temporarily deployed to test August 5, 2026 16:55 Inactive

chtruong814 removed the community-request label Aug 5, 2026

mathemakitten approved these changes Aug 5, 2026

View reviewed changes

svcnvidia-nemo-ci added the Final Review PR is in the "final review" stage label Aug 5, 2026

This was referenced Aug 5, 2026

fix: use explicit Megatron-FSDP V1/V2 classes in isinstance checks NVIDIA-NeMo/Megatron-Bridge#5357

Closed

fix(agent): make the functional-test report say what it actually tested NVIDIA-NeMo/RL#3524

Merged

santhnm2 added 2 commits August 6, 2026 13:03

Fix test and formatting

ea2189d

Signed-off-by: Keshav Santhanam <ksanthanam@nvidia.com>

Merge branch 'main' into ssm_inference_interface

83e1849

shanmugamr1992 approved these changes Aug 6, 2026

View reviewed changes

svcnvidia-nemo-ci added Approved All necessary approvals have been made and removed Final Review PR is in the "final review" stage labels Aug 6, 2026

santhnm2 added 2 commits August 7, 2026 10:54

Merge branch 'main' into ssm_inference_interface

d0c6654

Merge branch 'main' into ssm_inference_interface

8884612

copy-pr-bot Bot temporarily deployed to public August 7, 2026 19:25 Inactive

copy-pr-bot Bot temporarily deployed to test August 7, 2026 19:26 Inactive

copy-pr-bot Bot temporarily deployed to public August 7, 2026 19:29 Inactive

copy-pr-bot Bot temporarily deployed to public August 7, 2026 19:30 Inactive

copy-pr-bot Bot temporarily deployed to public August 7, 2026 19:40 Inactive

santhnm2 added this pull request to the merge queue Aug 8, 2026

Merged via the queue into NVIDIA:main with commit d12f6c8 Aug 8, 2026
183 checks passed

santhnm2 deleted the ssm_inference_interface branch August 8, 2026 21:05

santhnm2 mentioned this pull request Aug 11, 2026

Refactor GDP dynamic inference onto the shared SSM interface #6443

Merged

6 tasks

anlthms mentioned this pull request Aug 13, 2026

Add GDN dynamic inference #6499

Merged

6 tasks

Navigation Menu

Uh oh!

Add generic interface for SSM inference - #5382

Add generic interface for SSM inference#5382
santhnm2 merged 10 commits into
NVIDIA:mainfrom
santhnm2:ssm_inference_interface

santhnm2 commented Jun 16, 2026 •

edited

Loading

Uh oh!

copy-pr-bot Bot commented Jun 16, 2026

Uh oh!

santhnm2 commented Aug 4, 2026

Uh oh!

shanmugamr1992 commented Aug 5, 2026 •

edited

Loading

Uh oh!

santhnm2 commented Aug 5, 2026

Uh oh!

santhnm2 commented Aug 6, 2026

Uh oh!

santhnm2 commented Aug 7, 2026

Uh oh!

svcnvidia-nemo-ci commented Aug 8, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

Uh oh!

Conversation

santhnm2 commented Jun 16, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What does this PR do ?

Issue tracking

Contribution process

Pre-checks

Code review

Step 1: Mark PR as "Ready for Review"

Step 2: Final Review

Step 3: Approved

Merge

Uh oh!

copy-pr-bot Bot commented Jun 16, 2026

Uh oh!

santhnm2 commented Aug 4, 2026

Uh oh!

shanmugamr1992 commented Aug 5, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

NeMo-RL functional tests

Uh oh!

santhnm2 commented Aug 5, 2026

Uh oh!

santhnm2 commented Aug 6, 2026

Uh oh!

santhnm2 commented Aug 7, 2026

Uh oh!

svcnvidia-nemo-ci commented Aug 8, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

santhnm2 commented Jun 16, 2026 •

edited

Loading

shanmugamr1992 commented Aug 5, 2026 •

edited

Loading